from ultralytics import YOLO
import cv2
import json
import argparse
import math
import numpy as np
import traceback
import sys
import os
import re
import time

# 预编译正则表达式（性能优化）
chinese_pattern = re.compile('[\u4e00-\u9fff]')

# 全局变量保存模型
g_model = None

def parse_args():
    """统一解析命令行参数，同时支持交互和非交互模式"""
    parser = argparse.ArgumentParser(description="YOLOv8 Card Corner Detection")
    # 交互模式参数
    parser.add_argument("--interactive", action="store_true", help="Run in interactive mode")
    # 非交互模式参数（可选，仅在非交互模式下需要）
    parser.add_argument("--model_path", help="YOLOv8 .pt model path")
    parser.add_argument("--front_img", help="Front card image path")
    parser.add_argument("--back_img", help="Back card image path")
    return parser.parse_args()

def send_progress(percentage, message):
    """发送进度信息到Qt端"""
    progress = {
        "progress": percentage,
        "message": message
    }
    print(json.dumps(progress, ensure_ascii=False), flush=True)

def has_chinese_chars(s):
    """判断字符串是否包含中文字符"""
    return bool(chinese_pattern.search(s))

def decode_windows_path(path_str):
    """修复Windows系统下命令行参数的中文乱码问题"""
    try:
        if has_chinese_chars(path_str):
            return path_str
            
        # 尝试将可能乱码的字符串解码为正确中文
        return path_str.encode('latin-1').decode('gbk')
    except UnicodeEncodeError:
        return path_str
    except UnicodeDecodeError:
        try:
            return path_str.encode('latin-1').decode('utf-8', errors='replace')
        except:
            return path_str

def sort_corners(corners):
    """角点顺时针排序（左上→右上→右下→左下）"""
    if len(corners) != 4:
        return corners

    corners_np = np.array(corners)
    center_x = corners_np[:, 0].mean()
    center_y = corners_np[:, 1].mean()

    def angle(corner):
        x, y = corner
        return math.atan2(y - center_y, x - center_x)

    return sorted(corners, key=angle)

def infer_single_image(model, img_path):
    score_thresh = 0.25        # 与推理 conf 一致
    corner_conf_thresh = 0.25  # 角点置信度阈值

    # 读取图像（支持Windows中文路径）
    try:
        if not os.path.exists(img_path):
            raise FileNotFoundError(f"文件不存在: {img_path}")
            
        if not os.access(img_path, os.R_OK):
            raise PermissionError(f"没有文件读取权限: {img_path}")
            
        img_bytes = np.fromfile(img_path, dtype=np.uint8)
        if img_bytes.size == 0:
            raise ValueError(f"文件为空: {img_path}")
            
        img = cv2.imdecode(img_bytes, cv2.IMREAD_COLOR)
        if img is None:
            raise ValueError(f"图像解码失败: {img_path}（可能是损坏的图片或不支持的格式）")
            
    except Exception as e:
        raise ValueError(f"读取图像失败: {img_path}，错误: {str(e)}")

    # 推理（参数与测试代码一致）
    results = model(
        img,
        imgsz=640,
        conf=score_thresh,
        iou=0.7,
        rect=False,
        augment=False,
        verbose=False
    )

    infer_results = []
    for r in results:
        boxes = r.boxes.data.cpu().numpy()  # shape: (num_boxes, 6) → [x1,y1,x2,y2,conf,class]
        if len(boxes) == 0:
            continue

        keypoints = r.keypoints.data.cpu().numpy()
        if len(keypoints) == 0:
            continue

        if len(boxes) != len(keypoints):
            print(f"警告: 边界框数量 ({len(boxes)}) 与关键点数量 ({len(keypoints)}) 不匹配", file=sys.stderr)
            continue

        # 遍历每个检测结果
        for box, kpt in zip(boxes, keypoints):
            x1, y1, x2, y2, box_conf, _ = box
            bbox = [int(x1), int(y1), int(x2 - x1), int(y2 - y1)]

            if box_conf < score_thresh:
                continue

            # 过滤格式正确且置信度达标的关键点
            if kpt.shape[1] == 3:  # 确保关键点格式正确
                valid_mask = kpt[:, 2] > corner_conf_thresh  # 置信度掩码
                valid_kpts = kpt[valid_mask]  # 过滤后的关键点
                corners = valid_kpts[:, :2].astype(float).tolist()  # 提取x,y并转为列表
            else:
                print(f"警告: 关键点格式错误 (预期3个值，实际{ kpt.shape[1] }个)", file=sys.stderr)
                corners = []

            sorted_corners = sort_corners(corners)
            if len(sorted_corners) != 4:
                print(f"警告: 有效角点数量不足4个 (实际{len(sorted_corners)}个)，跳过", file=sys.stderr)
                continue

            infer_results.append({
                "bbox": bbox,
                "confidence": float(box_conf),
                "corners": sorted_corners
            })

    return infer_results

def process_command(command):
    """处理来自Qt的命令"""
    global g_model
    
    try:
        cmd = json.loads(command)
        
        # 处理加载模型命令
        if cmd["command"] == "load_model":
            model_path = decode_windows_path(cmd["model_path"])
            
            send_progress(50, "开始加载模型...")
            
            # 检查模型文件
            if not os.path.exists(model_path):
                return {
                    "status": "error",
                    "command": "load_model",
                    "message": f"模型文件不存在: {model_path}"
                }
                
            # 加载模型
            start_time = time.time()
            g_model = YOLO(model_path)
            load_time = time.time() - start_time
            
            if not g_model:
                return {
                    "status": "error",
                    "command": "load_model",
                    "message": f"模型加载失败: {model_path}"
                }
                
            send_progress(90, "模型加载完成，初始化中...")
            
            return {
                "status": "success",
                "command": "load_model",
                "message": f"模型加载成功，耗时{load_time:.2f}秒",
                "model_path": model_path
            }
            
        # 处理检测命令
        elif cmd["command"] == "detect":
            if g_model is None:
                return {
                    "status": "error",
                    "command": "detect",
                    "message": "模型未加载，请先加载模型"
                }
                
            front_img = decode_windows_path(cmd["front_img"])
            back_img = decode_windows_path(cmd["back_img"])
            
            send_progress(0, "开始检测正面图像...")
            front_results = infer_single_image(g_model, front_img)
            
            send_progress(50, "开始检测反面图像...")
            back_results = infer_single_image(g_model, back_img)
            
            return {
                "status": "success",
                "command": "detect",
                "front": front_results,
                "back": back_results
            }
            
        # 处理退出命令
        elif cmd["command"] == "exit":
            return {
                "status": "success",
                "command": "exit",
                "message": "程序即将退出"
            }
            
        # 未知命令
        else:
            return {
                "status": "error",
                "message": f"未知命令: {cmd['command']}"
            }
            
    except Exception as e:
        traceback.print_exc()
        return {
            "status": "error",
            "message": str(e),
            "traceback": traceback.format_exc()
        }

def run_interactive_mode():
    """运行交互模式，持续接收并处理命令"""
    global g_model
    
    send_progress(0, "进入交互模式，等待命令...")
    
    try:
        while True:
            # 读取Qt发送的命令（一行一个JSON）
            command = sys.stdin.readline()
            if not command:
                time.sleep(0.1)
                continue
                
            command = command.strip()
            if not command:
                continue
                
            # 处理特殊退出命令
            if command.lower() in ["exit", "quit"]:
                print(json.dumps({
                    "status": "success",
                    "message": "收到退出命令，正在退出..."
                }, ensure_ascii=False), flush=True)
                break
                
            # 处理常规命令
            result = process_command(command)
            print(json.dumps(result, ensure_ascii=False), flush=True)
            
            # 如果是退出命令，处理后结束循环
            if result.get("command") == "exit" and result.get("status") == "success":
                break
                
    except KeyboardInterrupt:
        print(json.dumps({
            "status": "error",
            "message": "收到中断信号"
        }, ensure_ascii=False), flush=True)
    except Exception as e:
        traceback.print_exc()
        print(json.dumps({
            "status": "error",
            "message": f"交互模式异常: {str(e)}",
            "traceback": traceback.format_exc()
        }, ensure_ascii=False), flush=True)
    finally:
        # 清理资源
        g_model = None
        send_progress(100, "程序已退出")

def run_single_mode(args):
    """运行单张图片处理模式"""
    try:
        # 验证必要参数
        if not args.model_path or not args.front_img or not args.back_img:
            raise ValueError("缺少必要参数: model_path, front_img, back_img 均为必填项")

        if sys.platform.startswith('win'):
            args.model_path = decode_windows_path(args.model_path)
            args.front_img = decode_windows_path(args.front_img)
            args.back_img = decode_windows_path(args.back_img)

        send_progress(30, "正在加载模型...")
        model = YOLO(args.model_path)
        if not model:
            raise ValueError(f"模型加载失败: {args.model_path}")
            
        send_progress(60, "正在检测图像...")
        front_results = infer_single_image(model, args.front_img)
        back_results = infer_single_image(model, args.back_img)

        final_result = {
            "status": "success",
            "front": front_results,
            "back": back_results
        }
        print(json.dumps(final_result, ensure_ascii=False), flush=True)

    except Exception as e:
        traceback.print_exc()
        error_result = {
            "status": "error",
            "message": str(e),
            "traceback": traceback.format_exc()
        }
        print(json.dumps(error_result, ensure_ascii=False), flush=True)
        exit(1)

if __name__ == "__main__":
    args = parse_args()
    
    # 根据参数选择运行模式
    if args.interactive:
        run_interactive_mode()
    else:
        run_single_mode(args)
